from utils import print_timing, is_blankline
try:
    import genocpp
except ImportError:
    print "import of genocpp module failed, attempting to compile..."
    import os
    os.system("scons")
    import genocpp

class Geno_Base(object):
    def __init__(self, genofile, schema, dbstring, tmpdir, table_dict, jobs):
        self.genofile = genofile
        self.schema = schema
        self.dbstring = dbstring
        self.tmpdir = tmpdir
        self.table_dict = table_dict
        self.jobs = jobs

    def get_idlink_ids(self, session):
        if 'idlink_ids' in self.__dict__:
            return self.idlink_ids
        else:
            Idlink = self.table_dict['Idlink']
            expids =  self.celfilenames()
            iquery = session.query(Idlink)
            self.idlink_ids = [iquery.filter(Idlink.expid==expid).one().id for expid in expids] # Query db for ids corresponding to expids
            return self.idlink_ids

    def truncate_genofile(self, truncgenofile, idlist):
        import genocpp, os, sys
        from utils import file_not_exist, safe_mkdir
        if [x for x in [self.genofile,] if file_not_exist(x)]:
            sys.exit("Exiting in Geno_Base.truncate_genofile in geno.py...")
        safe_mkdir(self.tmpdir)
        truncgenofile = os.path.join(self.tmpdir, truncgenofile)
        print "truncating geno file %s to %s"%(self.genofile, truncgenofile)
        discard_lines =  self.discard_lines_before_celfilenames() + 1

        linenum=0
        gf = open(self.genofile)
        tgf = open(truncgenofile, 'w')
        # Write header to trunc file.
        for i in range(discard_lines):
            tgf.write(gf.readline())
        for line in gf:
            linenum += 1
            line_split = line.lstrip().rstrip().split("\t")
            id = line_split[0]
            calls = line_split[1:]
            if id in idlist:
                tgf.write(line)
        gf.close()
        tgf.close()

        #genocpp.cpp_truncate_genofile(self.genofile, truncgenofile, idlist, discard_lines)

    def copy_data_to_geno_table(self, session, genotable):
        # Wrap TRUNCATE geno; COPY in COMMIT; BEGIN; -- COMMIT; to ensure that they run in the same transaction
        # If TRUNCATE is run in the same transaction as COPY, the write-ahead-logs (wal) will not be written to,
        # and the copy will progress faster.
        # See http://www.postgresql.org/docs/current/static/sql-begin.html
        import os
        session.execute("COMMIT; BEGIN; TRUNCATE %s.%s; COPY %s.%s FROM '%s' USING DELIMITERS ','; COMMIT;"%(self.schema, genotable, self.schema, genotable, self.genotablefile))

    def drop_geno_table_constraints(self, session, genotable):
        session.execute("select * from drop_constraint_if_exists('%s', '%s_idlink_id_fkey', '%s');"%(genotable, genotable, self.schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_anno_id_fkey', '%s');"%(genotable, genotable, self.schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_snpval_id_fkey', '%s');"%(genotable, genotable, self.schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_pkey', '%s');"%(genotable, genotable, self.schema))

    def restore_geno_idlink_id_fkey(self, session, genotable):
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_idlink_id_fkey FOREIGN KEY (idlink_id) REFERENCES %s.idlink(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, genotable, genotable, self.schema))

    def restore_geno_anno_id_fkey(self, session, genotable):
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_anno_id_fkey FOREIGN KEY (anno_id) REFERENCES %s.anno(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, genotable, genotable, self.schema))

    def restore_geno_snpval_id_fkey(self, session, genotable):
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_snpval_id_fkey FOREIGN KEY (snpval_id) REFERENCES %s.snpval(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, genotable, genotable, self.schema))

    def restore_geno_pkey(self, session, genotable):
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_pkey PRIMARY KEY (idlink_id, anno_id);"%(self.schema, genotable, genotable))

    def restore_geno_table_constraints(self, session, genotable):
        self.restore_geno_idlink_id_fkey(session, genotable)
        self.restore_geno_anno_id_fkey(session, genotable)
        self.restore_geno_snpval_id_fkey(session, genotable)
        self.restore_geno_pkey(session, genotable)

class Affy6(Geno_Base):

    def discard_lines_before_celfilenames(self):
        """Check for 'probeset_id' at beginning of expid line"""
        af = open(self.genofile)
        linenum=0
        while True:
            line = af.readline()
            if 'probeset_id' in line:
                break
            if not line or linenum > 10000:
                import sys
                print "Couldn't find string 'probeset_id' in the first %s lines of %s"%(linenum, self.genofile)
                print "Is it a correctly formatted Affymetrix 6 genotype file?"
                sys.exit(1)
            linenum = linenum + 1
        af.close()
        return linenum

    def discard_cols_before_celfilenames(self):
        discard_lines = self.discard_lines_before_celfilenames()
        discard_cols=0
        gf = open(self.genofile)
        for i in range(discard_lines):
            gf.readline()
        line_split = gf.readline().lstrip().rstrip().split("\t")
        gf.close()
        #Skip lines before first celfilename.
        l = len(line_split)
        for i, val in enumerate(line_split):
            if "CEL" in val:
                break
            if i == l-1:
                print "Error: reached the end of the experimental id line in %s without finding the string 'CEL'"%self.genofile
                print "Please check that the file is a correctly formatted Affymetrix 6 file"
                import sys
                sys.exit(1)
            discard_cols += 1
        return discard_cols

    def celfilenames(self):
        """
        Chip IDs: These correspond to the patients. In the case of
        Affymetrix these are Celfilenames. These all appear on one line in
        the calls file.
        """
        gf = open(self.genofile)
        for i in range(self.discard_lines_before_celfilenames()):
            gf.readline()
        fids = []
        discard_cols = self.discard_cols_before_celfilenames()
        line = gf.readline().lstrip().rstrip().split("\t")[discard_cols:]
        gf.close()
        import locale
        return line

class Illumina(Geno_Base):

    def discard_lines_before_celfilenames(self):
        """Check for '[Data]' on line preceding expid line"""
        gf = open(self.genofile)
        linenum=0
        while True:
            line = gf.readline()
            linenum = linenum + 1
            if '[Data]' in line:
                break
            if not line or linenum > 10000:
                import sys
                print "Couldn't find string '[Data]' in the first %s lines of %s"%(linenum, self.genofile)
                print "Is it a correctly formatted Illumina genotype file?"
                sys.exit(1)
        gf.close()
        return linenum

    # dummy function for compatability with tables_affy6.py
    def discard_cols_before_celfilenames(self):
        return 0

    def celfilenames(self):
        """
        Chip IDs: These correspond to the patients. In the case of
        Affymetrix these are Celfilenames. These all appear on one line in
        the calls file.
        """
        af = open(self.genofile)
        discard_lines = self.discard_lines_before_celfilenames()
        for i in range(discard_lines):
            af.readline()
        fids = []
        discard_cols = self.discard_cols_before_celfilenames()
        line = af.readline().lstrip().rstrip().split("\t")[discard_cols:]
        af.close()
        import locale
        return line

class Geno_Single(Geno_Base):
    def __init__(self, genofile, schema, dbstring, tmpdir, table_dict, jobs):
        import os
        Geno_Base.__init__(self, genofile, schema, dbstring, tmpdir, table_dict, jobs)
        self.genotablefile = os.path.abspath(os.path.join(self.tmpdir, "geno_%s"%self.schema))

    @print_timing
    def make_geno_from_file(self, session):
        import os, sys
        self.write_geno_to_file(session)
        session.execute("set search_path to %s"%self.schema)
        session.execute("ALTER TABLE %s.geno SET ( autovacuum_enabled = false );"%self.schema)
        session.commit()
        self.drop_geno_table_constraints(session, 'geno')
        self.copy_data_to_geno_table(session, 'geno')
        self.restore_geno_table_constraints(session, 'geno')
        session.execute("ALTER TABLE %s.geno SET ( autovacuum_enabled = true );"%self.schema)
        session.close()

    @print_timing
    def make_geno(self):
        import os, shutil, sys
        from utils import safe_mkdir
        safe_mkdir(self.tmpdir)
        from utils import create_engine_wrapper
        from sqlalchemy.orm import mapper, relation, sessionmaker
        import logging
        logging.basicConfig()
        db = create_engine_wrapper(self.dbstring)
        meta = self.table_dict['metadata']
        meta.bind = db
        meta.create_all()
        Session = sessionmaker(bind=db)
        session = Session()
        # Check that anno table not empty.
        result = session.execute("SELECT * FROM %s.anno;"%self.schema).fetchall()
        session.commit()
        if len(result) == 0:
            sys.exit("anno table is empty, exiting")
        session.close()
        self.make_geno_from_file(session)
        if os.path.isdir(self.tmpdir):
            shutil.rmtree(self.tmpdir)

    @print_timing
    def copy_data_to_geno_table(self, session, genotable):
        print "writing csv file %s to geno table..."%self.genotablefile
        super(Geno_Single, self).copy_data_to_geno_table(session, genotable)

    @print_timing
    def drop_geno_table_constraints(self, session, genotable):
        print "dropping geno table constraints..."
        super(Geno_Single, self).drop_geno_table_constraints(session, genotable)

    @print_timing
    def restore_geno_idlink_id_fkey(self, session, genotable):
        super(Geno_Single, self).restore_geno_idlink_id_fkey(session, genotable)

    @print_timing
    def restore_geno_anno_id_fkey(self, session, genotable):
        super(Geno_Single, self).restore_geno_anno_id_fkey(session, genotable)

    @print_timing
    def restore_geno_snpval_id_fkey(self, session, genotable):
        super(Geno_Single, self).restore_geno_snpval_id_fkey(session, genotable)

    @print_timing
    def restore_geno_pkey(self, session, genotable):
        super(Geno_Single, self).restore_geno_pkey(session, genotable)

    @print_timing
    def restore_geno_table_constraints(self, session, genotable):
        super(Geno_Single, self).restore_geno_table_constraints(session, genotable)

class Geno_Single_Affy6(Geno_Single, Affy6):

    @print_timing
    def write_geno_to_file(self, session):
        import os, shutil
        # rows of calls file correspond to snps.
        # columns of calls file correspond to patients.
        print "writing csv file %s containing geno table data..."%self.genotablefile
        from sqlalchemy.sql import select
        from utils import get_conf
        td = self.table_dict
        Idlink, idlink_table, Anno, anno_table = td['Idlink'], td['idlink_table'], td['Anno'], td['anno_table']
        # Discard lines before celfilenames + 1 for celfilename header.
        discard_lines =  self.discard_lines_before_celfilenames() + 1
        discard_cols = self.discard_cols_before_celfilenames()
        idlink_ids = self.get_idlink_ids(session)
        anno_stmt = select([anno_table.c.id, anno_table.c.fid])
        annodict = {}
        for row in anno_stmt.execute():
            annodict[row[1]] = int(row[0])
        conf = get_conf()
        sync_interval = int(conf["sync_interval"])
        session.commit()
        session.flush()
        genocpp.write_geno_table_affy6_to_file(discard_lines, discard_cols, sync_interval, self.genofile, self.genotablefile, idlink_ids, annodict)

class Geno_Single_Illumina(Geno_Single, Illumina):
    def __init__(self, genofile, genotablefile, schema, dbstring, table_dict, alleletype, jobs):
        Geno_Single.__init__(self, genofile, genotablefile, schema, dbstring, table_dict, jobs)
        self.alleletype = alleletype

    @print_timing
    def write_geno_to_file(self, session):
        # rows of calls file correspond to snps.
        # columns of calls file correspond to patients.
        print "writing csv file %s containing geno table data..."%self.genotablefile
        from sqlalchemy.sql import select
        from utils import get_conf
        td = self.table_dict
        anno_table, geno_table, Idlink, Anno = td['anno_table'], td['geno_table'], td['Idlink'], td['Anno']
        # Discard lines before celfilenames + 1 for celfilename header.
        discard_lines =  self.discard_lines_before_celfilenames() + 1
        linenum=0
        idlink_ids = self.get_idlink_ids(session)
        stmt = select([anno_table.c.rsid, anno_table.c.id, anno_table.c.allelea_forward_id, anno_table.c.alleleb_forward_id, anno_table.c.allelea_top_id, anno_table.c.alleleb_top_id])
        annodict = {}
        error_rsid = set()
        for row in stmt.execute():
            annodict[row[0]] = {'id':str(row[1]), 'allelea_forward':row[2], 'alleleb_forward':row[3], 'allelea_top':row[4], 'alleleb_top':row[5]}
        conf = get_conf()
        sync_interval = int(conf["sync_interval"])
        genocpp.write_geno_table_illumina_to_file(discard_lines, sync_interval, self.genofile, self.genotablefile, idlink_ids, annodict, self.alleletype)

# This function is used for multiprocessing, so can't be a member function of a class.
#def load_geno_table_from_file(dbstring, genotable, genotablefile, schema, i):
def load_geno_table_from_file(arglist):
    dbstring, genotable, genotablefile, schema, i = arglist['dbstring'], arglist['genotable'], arglist['genotablefile'], arglist['schema'], arglist['i']
    import os
    from utils import create_engine_wrapper
    from sqlalchemy.orm import sessionmaker
    db = create_engine_wrapper(dbstring)
    Session = sessionmaker()
    session = Session(bind=db)
    try:
        session.execute("ALTER TABLE %s.geno%s SET ( autovacuum_enabled = false );"%(schema, i))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_idlink_id_check', '%s');"%(genotable, genotable, schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_idlink_id_fkey', '%s');"%(genotable, genotable, schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_anno_id_fkey', '%s');"%(genotable, genotable, schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_snpval_id_fkey', '%s');"%(genotable, genotable, schema))
        session.execute("select * from drop_constraint_if_exists('%s', '%s_pkey', '%s');"%(genotable, genotable, schema))
        session.execute("COMMIT; BEGIN; TRUNCATE %s.%s; COPY %s.%s FROM '%s' USING DELIMITERS ','; COMMIT;"%(schema, genotable, schema, genotable, genotablefile))
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_idlink_id_check CHECK ((idlink_id = %s));"%(schema, genotable, genotable, genotable.split('geno')[1]))
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_idlink_id_fkey FOREIGN KEY (idlink_id) REFERENCES %s.idlink(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(schema, genotable, genotable, schema))
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_anno_id_fkey FOREIGN KEY (anno_id) REFERENCES %s.anno(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(schema, genotable, genotable, schema))
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_snpval_id_fkey FOREIGN KEY (snpval_id) REFERENCES %s.snpval(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(schema, genotable, genotable, schema))
        session.execute("ALTER TABLE ONLY %s.%s ADD CONSTRAINT %s_pkey PRIMARY KEY (idlink_id, anno_id);"%(schema, genotable, genotable))
        session.execute("ALTER TABLE %s.geno%s SET ( autovacuum_enabled = true );"%(schema, i))
    except Exception, e:
        # Work around Python bug http://bugs.python.org/issue1692335 See also
        # http://stackoverflow.com/questions/8785899/hang-in-python-script-using-sqlalchemy-and-multiprocessing
        import sys
        raise Exception(e.__class__.__name__ + ": " +str(e)), None, sys.exc_info()[2]
    finally:
        session.commit()
        session.close()
        db.dispose()

class Geno_Shard(Geno_Base):
    @print_timing
    def make_geno(self):
        import os, shutil, sys
        from utils import safe_mkdir
        safe_mkdir(self.tmpdir)
        from sqlalchemy.sql import select
        idlink_table = self.table_dict['idlink_table']
        from utils import create_engine_wrapper
        from sqlalchemy.orm import mapper, relation, sessionmaker
        meta = self.table_dict['metadata']
        db = create_engine_wrapper(self.dbstring)
        #db.echo = 'debug'
        meta.bind = db
        meta.create_all()
        Session = sessionmaker(bind=db)
        session = Session()
        ids = self.get_idlink_ids(session)
        import logging
        logging.basicConfig()
        # Check that anno table not empty.
        result = session.execute("SELECT * FROM %s.anno;"%self.schema).fetchall()
        session.commit()
        if len(result) == 0:
            sys.exit("anno table is empty, exiting")
        session.execute("DROP TABLE IF EXISTS %s.geno_idlink CASCADE"%self.schema)
        for i in ids:
            session.execute('DROP TABLE IF EXISTS %s.geno%s'%(self.schema, i))
        session.commit()
        self.make_geno_from_file(session)
        if os.path.isdir(self.tmpdir):
            shutil.rmtree(self.tmpdir)

    @print_timing
    def make_geno_from_file(self, session):
        import os, sys
        from sqlalchemy.sql import select
        from sqlalchemy.orm import sessionmaker
        from dbschema import make_geno_idlink_table, make_geno_shard_table
        meta = self.table_dict['metadata']
        Idlink = self.table_dict['Idlink']
        idlink_table = self.table_dict['idlink_table']
        session.execute("set search_path to %s"%self.schema)
        ids = self.get_idlink_ids(session)
        for i in ids:
            if os.path.exists('geno%s'%i):
                os.remove('geno%s'%i)
        geno_table_dict = {}
        self.table_dict['geno_idlink_table'] = make_geno_idlink_table(meta, self.schema)
        for i in ids:
            geno_table_dict[i] = make_geno_shard_table(meta, self.schema, i)
        session.commit()
        meta.create_all()
        session.commit()
        self.write_geno_to_file(session)
        self.load_geno_idlink_table(session)
        session.close()
        self.make_geno_tables_from_file_threaded()

    def load_geno_idlink_table(self, session):
        td = self.table_dict
        geno_idlink_table, Idlink = td['geno_idlink_table'], td['Idlink']
        insert_dict_list = []
        ids = self.get_idlink_ids(session)
        for i in ids:
            insert_dict_list.append({'id':i})
        session.execute(geno_idlink_table.insert(), insert_dict_list)
        session.commit()

    @print_timing
    def make_geno_tables_from_file_threaded(self):
        import os
        from utils import create_engine_wrapper
        from sqlalchemy.orm import scoped_session, sessionmaker
        from multiprocessing import Pool
        db = create_engine_wrapper(self.dbstring)
        meta = self.table_dict['metadata']
        meta.bind = db
        meta.create_all()
        Session = sessionmaker()
        session = Session(bind=db)
        pool = Pool(processes=self.jobs)
        ids = self.get_idlink_ids(session)
        arglist = []
        for i in ids:
            arglist.append({'dbstring':self.dbstring, 'genotable':'geno%s'%i, 'genotablefile':os.path.abspath(os.path.join(self.tmpdir, 'geno%s'%i)), 'schema':self.schema, 'i':i})
        results = []
        r = pool.map_async(load_geno_table_from_file, arglist, callback=results.append)
        r.get()
        r.wait()
        pool.close()
        pool.join()

class Geno_Shard_Affy6(Geno_Shard, Affy6):
    @print_timing
    def write_geno_to_file(self, session):
        import os, shutil
        # rows of calls file correspond to snps.
        # columns of calls file correspond to patients.
        print "writing csv files containing geno sharding table data..."
        from sqlalchemy.sql import select
        from utils import get_conf
        td = self.table_dict
        Idlink, idlink_table, Anno, anno_table = td['Idlink'], td['idlink_table'], td['Anno'], td['anno_table']
        # Discard lines before celfilenames + 1 for celfilename header.
        discard_lines =  self.discard_lines_before_celfilenames() + 1
        discard_cols = self.discard_cols_before_celfilenames()
        idlink_ids = self.get_idlink_ids(session)
        anno_stmt = select([anno_table.c.id, anno_table.c.fid])
        annodict = {}
        for row in anno_stmt.execute():
            annodict[row[1]] = int(row[0])
        conf = get_conf()
        sync_interval = int(conf["sync_interval"])
        session.commit()
        session.flush()
        genocpp.write_geno_tables_affy6_to_file(self.tmpdir, discard_lines, sync_interval, self.genofile, idlink_ids, annodict)

class Geno_Shard_Illumina(Geno_Shard, Illumina):
    def __init__(self, genofile, schema, dbstring, tmpdir, table_dict, alleletype, jobs):
        Geno_Shard.__init__(self, genofile, schema, dbstring, tmpdir, table_dict, jobs)
        self.alleletype = alleletype

    @print_timing
    def write_geno_to_file(self, session):
        # rows of calls file correspond to snps.
        # columns of calls file correspond to patients.
        print "writing csv files containing geno sharding table data..."
        from sqlalchemy.sql import select
        from utils import get_conf
        import os, shutil
        td = self.table_dict
        anno_table, idlink_table, Idlink, Anno = td['anno_table'], td['idlink_table'], td['Idlink'], td['Anno']
        # Discard lines before celfilenames + 1 for celfilename header.
        discard_lines =  self.discard_lines_before_celfilenames() + 1
        linenum=0
        gtdict = {}
        idlink_ids = self.get_idlink_ids(session)
        idlen = len(idlink_ids)
        gfdict = {}
        iquery = session.query(Idlink)
        anno_stmt = select([anno_table.c.rsid, anno_table.c.id, anno_table.c.allelea_forward_id, anno_table.c.alleleb_forward_id, anno_table.c.allelea_top_id, anno_table.c.alleleb_top_id])
        annodict = {}
        error_rsid = set()
        for row in anno_stmt.execute():
            annodict[row[0]] = {'id':str(row[1]), 'allelea_forward':row[2], 'alleleb_forward':row[3], 'allelea_top':row[4], 'alleleb_top':row[5]}
        expids = self.celfilenames()
        conf = get_conf()
        sync_interval = int(conf["sync_interval"])
        genocpp.write_geno_tables_illumina_to_file(self.tmpdir, discard_lines, sync_interval, self.genofile, idlink_ids, annodict, self.alleletype)
